# -*- coding: utf-8 -*-
"""
@Time ： 2024/4/3 9:06
@Auth ： fcq
@File ：BERT_SentiX.py
@IDE ：PyCharm
@Motto：ABC(Always Be Coding)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class BERT_SentiX(nn.Module):
    def __init__(self,  opt,bert,):
        super(BERT_SentiX, self).__init__()
        self.bert = bert
        self.dropout = nn.Dropout(opt.dropout)
        self.dense = nn.Linear(opt.bert_dim, opt.num_labels)

    def forward(self, inputs):
        text_bert_indices, bert_segments_ids, attention_mask, sentix_input = inputs[0], inputs[1], inputs[2], inputs[3]



        pooled_output = self.bert(text_bert_indices, token_type_ids=bert_segments_ids, attention_mask=attention_mask)
        pooled_output = self.dropout(pooled_output['pooler_output'])

        feature = pooled_output.unsqueeze(1)
        feature = F.normalize(feature, dim=2)

        logits = self.dense(pooled_output)

        return logits, feature